{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from config import (BATCH_SIZE, CLIP_REWARD, DISCOUNT_FACTOR, ENV_NAME,\n",
    "                    EVAL_LENGTH, FRAMES_BETWEEN_EVAL, INPUT_SHAPE,\n",
    "                    LEARNING_RATE, LOAD_FROM, LOAD_REPLAY_BUFFER,\n",
    "                    MAX_EPISODE_LENGTH, MAX_NOOP_STEPS, MEM_SIZE,\n",
    "                    MIN_REPLAY_BUFFER_SIZE, PRIORITY_SCALE, SAVE_PATH,\n",
    "                    TENSORBOARD_DIR, TOTAL_FRAMES, UPDATE_FREQ, USE_PER,\n",
    "                    WRITE_TENSORBOARD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import cv2\n",
    "\n",
    "import random\n",
    "import os\n",
    "import json\n",
    "import time\n",
    "\n",
    "import gym\n",
    "\n",
    "import tensorflow as tf\n",
    "from tensorflow.keras.initializers import VarianceScaling\n",
    "from tensorflow.keras.layers import (Add, Conv2D, Dense, Flatten, Input,\n",
    "                                     Lambda, Subtract)\n",
    "from tensorflow.keras.models import Model\n",
    "from tensorflow.keras.optimizers import Adam, RMSprop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This function can resize to any shape, but was built to resize to 84x84\n",
    "def process_frame(frame, shape=(84, 84)):\n",
    "    \"\"\"Preprocesses a 210x160x3 frame to 84x84x1 grayscale\n",
    "    Arguments:\n",
    "        frame: The frame to process.  Must have values ranging from 0-255\n",
    "    Returns:\n",
    "        The processed frame\n",
    "    \"\"\"\n",
    "    frame = frame.astype(np.uint8)  # cv2 requires np.uint8, other dtypes will not work\n",
    "\n",
    "    frame = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)\n",
    "    frame = frame[34:34+160, :160]  # crop image\n",
    "    frame = cv2.resize(frame, shape, interpolation=cv2.INTER_NEAREST)\n",
    "    frame = frame.reshape((*shape, 1))\n",
    "\n",
    "    return frame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_q_network(n_actions, learning_rate=0.00001, input_shape=(84, 84), history_length=4):\n",
    "    \"\"\"Builds a dueling DQN as a Keras model\n",
    "    Arguments:\n",
    "        n_actions: Number of possible action the agent can take\n",
    "        learning_rate: Learning rate\n",
    "        input_shape: Shape of the preprocessed frame the model sees\n",
    "        history_length: Number of historical frames the agent can see\n",
    "    Returns:\n",
    "        A compiled Keras model\n",
    "    \"\"\"\n",
    "    model_input = Input(shape=(input_shape[0], input_shape[1], history_length))\n",
    "    x = Lambda(lambda layer: layer / 255)(model_input)  # normalize by 255\n",
    "\n",
    "    x = Conv2D(32, (8, 8), strides=4, kernel_initializer=VarianceScaling(scale=2.), activation='relu', use_bias=False)(x)\n",
    "    x = Conv2D(64, (4, 4), strides=2, kernel_initializer=VarianceScaling(scale=2.), activation='relu', use_bias=False)(x)\n",
    "    x = Conv2D(64, (3, 3), strides=1, kernel_initializer=VarianceScaling(scale=2.), activation='relu', use_bias=False)(x)\n",
    "    x = Conv2D(1024, (7, 7), strides=1, kernel_initializer=VarianceScaling(scale=2.), activation='relu', use_bias=False)(x)\n",
    "\n",
    "    # Split into value and advantage streams\n",
    "    val_stream, adv_stream = Lambda(lambda w: tf.split(w, 2, 3))(x)  # custom splitting layer\n",
    "\n",
    "    val_stream = Flatten()(val_stream)\n",
    "    val = Dense(1, kernel_initializer=VarianceScaling(scale=2.))(val_stream)\n",
    "\n",
    "    adv_stream = Flatten()(adv_stream)\n",
    "    adv = Dense(n_actions, kernel_initializer=VarianceScaling(scale=2.))(adv_stream)\n",
    "\n",
    "    # Combine streams into Q-Values\n",
    "    reduce_mean = Lambda(lambda w: tf.reduce_mean(w, axis=1, keepdims=True))  # custom layer for reduce mean\n",
    "\n",
    "    q_vals = Add()([val, Subtract()([adv, reduce_mean(adv)])])\n",
    "\n",
    "    # Build model\n",
    "    model = Model(model_input, q_vals)\n",
    "    model.compile(Adam(learning_rate), loss=tf.keras.losses.Huber())\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GameWrapper:\n",
    "    \"\"\"Wrapper for the environment provided by Gym\"\"\"\n",
    "    def __init__(self, env_name, no_op_steps=10, history_length=4):\n",
    "        self.env = gym.make(env_name)\n",
    "        self.no_op_steps = no_op_steps\n",
    "        self.history_length = 4\n",
    "\n",
    "        self.state = None\n",
    "        self.last_lives = 0\n",
    "\n",
    "    def reset(self, evaluation=False):\n",
    "        \"\"\"Resets the environment\n",
    "        Arguments:\n",
    "            evaluation: Set to True when the agent is being evaluated. Takes a random number of no-op steps if True.\n",
    "        \"\"\"\n",
    "\n",
    "        self.frame = self.env.reset()\n",
    "        self.last_lives = 0\n",
    "\n",
    "        # If evaluating, take a random number of no-op steps.\n",
    "        # This adds an element of randomness, so that the each\n",
    "        # evaluation is slightly different.\n",
    "        if evaluation:\n",
    "            for _ in range(random.randint(0, self.no_op_steps)):\n",
    "                self.env.step(1)\n",
    "\n",
    "        # For the initial state, we stack the first frame four times\n",
    "        self.state = np.repeat(process_frame(self.frame), self.history_length, axis=2)\n",
    "\n",
    "    def step(self, action, render_mode=None):\n",
    "        \"\"\"Performs an action and observes the result\n",
    "        Arguments:\n",
    "            action: An integer describe action the agent chose\n",
    "            render_mode: None doesn't render anything, 'human' renders the screen in a new window, 'rgb_array' returns an np.array with rgb values\n",
    "        Returns:\n",
    "            processed_frame: The processed new frame as a result of that action\n",
    "            reward: The reward for taking that action\n",
    "            terminal: Whether the game has ended\n",
    "            life_lost: Whether a life has been lost\n",
    "            new_frame: The raw new frame as a result of that action\n",
    "            If render_mode is set to 'rgb_array' this also returns the rendered rgb_array\n",
    "        \"\"\"\n",
    "        new_frame, reward, terminal, info = self.env.step(action)\n",
    "\n",
    "        # In the commonly ignored 'info' or 'meta' data returned by env.step\n",
    "        # we can get information such as the number of lives the agent has.\n",
    "\n",
    "        # We use this here to find out when the agent loses a life, and\n",
    "        # if so, we set life_lost to True.\n",
    "\n",
    "        # We use life_lost to force the agent to start the game\n",
    "        # and not sit around doing nothing.\n",
    "        if info['ale.lives'] < self.last_lives:\n",
    "            life_lost = True\n",
    "        else:\n",
    "            life_lost = terminal\n",
    "        self.last_lives = info['ale.lives']\n",
    "\n",
    "        processed_frame = process_frame(new_frame)\n",
    "        self.state = np.append(self.state[:, :, 1:], processed_frame, axis=2)\n",
    "\n",
    "        if render_mode == 'rgb_array':\n",
    "            return processed_frame, reward, terminal, life_lost, self.env.render(render_mode)\n",
    "        elif render_mode == 'human':\n",
    "            self.env.render()\n",
    "\n",
    "        return processed_frame, reward, terminal, life_lost"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ReplayBuffer:\n",
    "    \"\"\"Replay Buffer to store transitions.\n",
    "    This implementation was heavily inspired by Fabio M. Graetz's replay buffer\n",
    "    here: https://github.com/fg91/Deep-Q-Learning/blob/master/DQN.ipynb\"\"\"\n",
    "    def __init__(self, size=1000000, input_shape=(84, 84), history_length=4, use_per=True):\n",
    "        \"\"\"\n",
    "        Arguments:\n",
    "            size: Integer, Number of stored transitions\n",
    "            input_shape: Shape of the preprocessed frame\n",
    "            history_length: Integer, Number of frames stacked together to create a state for the agent\n",
    "            use_per: Use PER instead of classic experience replay\n",
    "        \"\"\"\n",
    "        self.size = size\n",
    "        self.input_shape = input_shape\n",
    "        self.history_length = history_length\n",
    "        self.count = 0  # total index of memory written to, always less than self.size\n",
    "        self.current = 0  # index to write to\n",
    "\n",
    "        # Pre-allocate memory\n",
    "        self.actions = np.empty(self.size, dtype=np.int32)\n",
    "        self.rewards = np.empty(self.size, dtype=np.float32)\n",
    "        self.frames = np.empty((self.size, self.input_shape[0], self.input_shape[1]), dtype=np.uint8)\n",
    "        self.terminal_flags = np.empty(self.size, dtype=np.bool)\n",
    "        self.priorities = np.zeros(self.size, dtype=np.float32)\n",
    "\n",
    "        self.use_per = use_per\n",
    "\n",
    "    def add_experience(self, action, frame, reward, terminal, clip_reward=True):\n",
    "        \"\"\"Saves a transition to the replay buffer\n",
    "        Arguments:\n",
    "            action: An integer between 0 and env.action_space.n - 1 \n",
    "                determining the action the agent perfomed\n",
    "            frame: A (84, 84, 1) frame of the game in grayscale\n",
    "            reward: A float determining the reward the agend received for performing an action\n",
    "            terminal: A bool stating whether the episode terminated\n",
    "        \"\"\"\n",
    "        if frame.shape != self.input_shape:\n",
    "            raise ValueError('Dimension of frame is wrong!')\n",
    "\n",
    "        if clip_reward:\n",
    "            reward = np.sign(reward)\n",
    "\n",
    "        # Write memory\n",
    "        self.actions[self.current] = action\n",
    "        self.frames[self.current, ...] = frame\n",
    "        self.rewards[self.current] = reward\n",
    "        self.terminal_flags[self.current] = terminal\n",
    "        self.priorities[self.current] = max(self.priorities.max(), 1)  # make the most recent experience important\n",
    "        self.count = max(self.count, self.current+1)\n",
    "        self.current = (self.current + 1) % self.size\n",
    "\n",
    "    def get_minibatch(self, batch_size=32, priority_scale=0.0):\n",
    "        \"\"\"Returns a minibatch of self.batch_size = 32 transitions\n",
    "        Arguments:\n",
    "            batch_size: How many samples to return\n",
    "            priority_scale: How much to weight priorities. 0 = completely random, 1 = completely based on priority\n",
    "        Returns:\n",
    "            A tuple of states, actions, rewards, new_states, and terminals\n",
    "            If use_per is True:\n",
    "                An array describing the importance of transition. Used for scaling gradient steps.\n",
    "                An array of each index that was sampled\n",
    "        \"\"\"\n",
    "\n",
    "        if self.count < self.history_length:\n",
    "            raise ValueError('Not enough memories to get a minibatch')\n",
    "\n",
    "        # Get sampling probabilities from priority list\n",
    "        if self.use_per:\n",
    "            scaled_priorities = self.priorities[self.history_length:self.count-1] ** priority_scale\n",
    "            sample_probabilities = scaled_priorities / sum(scaled_priorities)\n",
    "\n",
    "        # Get a list of valid indices\n",
    "        indices = []\n",
    "        for i in range(batch_size):\n",
    "            while True:\n",
    "                # Get a random number from history_length to maximum frame written with probabilities based on priority weights\n",
    "                if self.use_per:\n",
    "                    index = np.random.choice(np.arange(self.history_length, self.count-1), p=sample_probabilities)\n",
    "                else:\n",
    "                    index = random.randint(self.history_length, self.count - 1)\n",
    "\n",
    "                # We check that all frames are from same episode with the two following if statements.  If either are True, the index is invalid.\n",
    "                if index >= self.current and index - self.history_length <= self.current:\n",
    "                    continue\n",
    "                if self.terminal_flags[index - self.history_length:index].any():\n",
    "                    continue\n",
    "                break\n",
    "            indices.append(index)\n",
    "\n",
    "        # Retrieve states from memory\n",
    "        states = []\n",
    "        new_states = []\n",
    "        for idx in indices:\n",
    "            states.append(self.frames[idx-self.history_length:idx, ...])\n",
    "            new_states.append(self.frames[idx-self.history_length+1:idx+1, ...])\n",
    "\n",
    "        states = np.transpose(np.asarray(states), axes=(0, 2, 3, 1))\n",
    "        new_states = np.transpose(np.asarray(new_states), axes=(0, 2, 3, 1))\n",
    "\n",
    "        if self.use_per:\n",
    "            # Get importance weights from probabilities calculated earlier\n",
    "            importance = 1/self.count * 1/sample_probabilities[[index - self.history_length for index in indices]]\n",
    "            importance = importance / importance.max()\n",
    "\n",
    "            return (states, self.actions[indices], self.rewards[indices], new_states, self.terminal_flags[indices]), importance, indices\n",
    "        else:\n",
    "            return states, self.actions[indices], self.rewards[indices], new_states, self.terminal_flags[indices]\n",
    "\n",
    "    def set_priorities(self, indices, errors, offset=0.1):\n",
    "        \"\"\"Update priorities for PER\n",
    "        Arguments:\n",
    "            indices: Indices to update\n",
    "            errors: For each index, the error between the target Q-vals and the predicted Q-vals\n",
    "        \"\"\"\n",
    "        for i, e in zip(indices, errors):\n",
    "            self.priorities[i] = abs(e) + offset\n",
    "\n",
    "    def save(self, folder_name):\n",
    "        \"\"\"Save the replay buffer to a folder\"\"\"\n",
    "\n",
    "        if not os.path.isdir(folder_name):\n",
    "            os.mkdir(folder_name)\n",
    "\n",
    "        np.save(folder_name + '/actions.npy', self.actions)\n",
    "        np.save(folder_name + '/frames.npy', self.frames)\n",
    "        np.save(folder_name + '/rewards.npy', self.rewards)\n",
    "        np.save(folder_name + '/terminal_flags.npy', self.terminal_flags)\n",
    "\n",
    "    def load(self, folder_name):\n",
    "        \"\"\"Loads the replay buffer from a folder\"\"\"\n",
    "        self.actions = np.load(folder_name + '/actions.npy')\n",
    "        self.frames = np.load(folder_name + '/frames.npy')\n",
    "        self.rewards = np.load(folder_name + '/rewards.npy')\n",
    "        self.terminal_flags = np.load(folder_name + '/terminal_flags.npy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Agent(object):\n",
    "    \"\"\"Implements a standard DDDQN agent\"\"\"\n",
    "    def __init__(self,\n",
    "                 dqn,\n",
    "                 target_dqn,\n",
    "                 replay_buffer,\n",
    "                 n_actions,\n",
    "                 input_shape=(84, 84),\n",
    "                 batch_size=32,\n",
    "                 history_length=4,\n",
    "                 eps_initial=1,\n",
    "                 eps_final=0.1,\n",
    "                 eps_final_frame=0.01,\n",
    "                 eps_evaluation=0.0,\n",
    "                 eps_annealing_frames=1000000,\n",
    "                 replay_buffer_start_size=50000,\n",
    "                 max_frames=25000000,\n",
    "                 use_per=True):\n",
    "        \"\"\"\n",
    "        Arguments:\n",
    "            dqn: A DQN (returned by the DQN function) to predict moves\n",
    "            target_dqn: A DQN (returned by the DQN function) to predict target-q values.  This can be initialized in the same way as the dqn argument\n",
    "            replay_buffer: A ReplayBuffer object for holding all previous experiences\n",
    "            n_actions: Number of possible actions for the given environment\n",
    "            input_shape: Tuple/list describing the shape of the pre-processed environment\n",
    "            batch_size: Number of samples to draw from the replay memory every updating session\n",
    "            history_length: Number of historical frames available to the agent\n",
    "            eps_initial: Initial epsilon value.\n",
    "            eps_final: The \"half-way\" epsilon value.  The epsilon value decreases more slowly after this\n",
    "            eps_final_frame: The final epsilon value\n",
    "            eps_evaluation: The epsilon value used during evaluation\n",
    "            eps_annealing_frames: Number of frames during which epsilon will be annealed to eps_final, then eps_final_frame\n",
    "            replay_buffer_start_size: Size of replay buffer before beginning to learn (after this many frames, epsilon is decreased more slowly)\n",
    "            max_frames: Number of total frames the agent will be trained for\n",
    "            use_per: Use PER instead of classic experience replay\n",
    "        \"\"\"\n",
    "\n",
    "        self.n_actions = n_actions\n",
    "        self.input_shape = input_shape\n",
    "        self.history_length = history_length\n",
    "\n",
    "        # Memory information\n",
    "        self.replay_buffer_start_size = replay_buffer_start_size\n",
    "        self.max_frames = max_frames\n",
    "        self.batch_size = batch_size\n",
    "\n",
    "        self.replay_buffer = replay_buffer\n",
    "        self.use_per = use_per\n",
    "\n",
    "        # Epsilon information\n",
    "        self.eps_initial = eps_initial\n",
    "        self.eps_final = eps_final\n",
    "        self.eps_final_frame = eps_final_frame\n",
    "        self.eps_evaluation = eps_evaluation\n",
    "        self.eps_annealing_frames = eps_annealing_frames\n",
    "\n",
    "        # Slopes and intercepts for exploration decrease\n",
    "        # (Credit to Fabio M. Graetz for this and calculating epsilon based on frame number)\n",
    "        self.slope = -(self.eps_initial - self.eps_final) / self.eps_annealing_frames\n",
    "        self.intercept = self.eps_initial - self.slope*self.replay_buffer_start_size\n",
    "        self.slope_2 = -(self.eps_final - self.eps_final_frame) / (self.max_frames - self.eps_annealing_frames - self.replay_buffer_start_size)\n",
    "        self.intercept_2 = self.eps_final_frame - self.slope_2*self.max_frames\n",
    "\n",
    "        # DQN\n",
    "        self.DQN = dqn\n",
    "        self.target_dqn = target_dqn\n",
    "\n",
    "    def calc_epsilon(self, frame_number, evaluation=False):\n",
    "        \"\"\"Get the appropriate epsilon value from a given frame number\n",
    "        Arguments:\n",
    "            frame_number: Global frame number (used for epsilon)\n",
    "            evaluation: True if the model is evaluating, False otherwise (uses eps_evaluation instead of default epsilon value)\n",
    "        Returns:\n",
    "            The appropriate epsilon value\n",
    "        \"\"\"\n",
    "        if evaluation:\n",
    "            return self.eps_evaluation\n",
    "        elif frame_number < self.replay_buffer_start_size:\n",
    "            return self.eps_initial\n",
    "        elif frame_number >= self.replay_buffer_start_size and frame_number < self.replay_buffer_start_size + self.eps_annealing_frames:\n",
    "            return self.slope*frame_number + self.intercept\n",
    "        elif frame_number >= self.replay_buffer_start_size + self.eps_annealing_frames:\n",
    "            return self.slope_2*frame_number + self.intercept_2\n",
    "\n",
    "    def get_action(self, frame_number, state, evaluation=False):\n",
    "        \"\"\"Query the DQN for an action given a state\n",
    "        Arguments:\n",
    "            frame_number: Global frame number (used for epsilon)\n",
    "            state: State to give an action for\n",
    "            evaluation: True if the model is evaluating, False otherwise (uses eps_evaluation instead of default epsilon value)\n",
    "        Returns:\n",
    "            An integer as the predicted move\n",
    "        \"\"\"\n",
    "\n",
    "        # Calculate epsilon based on the frame number\n",
    "        eps = self.calc_epsilon(frame_number, evaluation)\n",
    "\n",
    "        # With chance epsilon, take a random action\n",
    "        if np.random.rand(1) < eps:\n",
    "            return np.random.randint(0, self.n_actions)\n",
    "\n",
    "        # Otherwise, query the DQN for an action\n",
    "        q_vals = self.DQN.predict(state.reshape((-1, self.input_shape[0], self.input_shape[1], self.history_length)))[0]\n",
    "        return q_vals.argmax()\n",
    "\n",
    "    def get_intermediate_representation(self, state, layer_names=None, stack_state=True):\n",
    "        \"\"\"\n",
    "        Get the output of a hidden layer inside the model.  This will be/is used for visualizing model\n",
    "        Arguments:\n",
    "            state: The input to the model to get outputs for hidden layers from\n",
    "            layer_names: Names of the layers to get outputs from.  This can be a list of multiple names, or a single name\n",
    "            stack_state: Stack `state` four times so the model can take input on a single (84, 84, 1) frame\n",
    "        Returns:\n",
    "            Outputs to the hidden layers specified, in the order they were specified.\n",
    "        \"\"\"\n",
    "        # Prepare list of layers\n",
    "        if isinstance(layer_names, list) or isinstance(layer_names, tuple):\n",
    "            layers = [self.DQN.get_layer(name=layer_name).output for layer_name in layer_names]\n",
    "        else:\n",
    "            layers = self.DQN.get_layer(name=layer_names).output\n",
    "\n",
    "        # Model for getting intermediate output\n",
    "        temp_model = tf.keras.Model(self.DQN.inputs, layers)\n",
    "\n",
    "        # Stack state 4 times\n",
    "        if stack_state:\n",
    "            if len(state.shape) == 2:\n",
    "                state = state[:, :, np.newaxis]\n",
    "            state = np.repeat(state, self.history_length, axis=2)\n",
    "\n",
    "        # Put it all together\n",
    "        return temp_model.predict(state.reshape((-1, self.input_shape[0], self.input_shape[1], self.history_length)))\n",
    "\n",
    "    def update_target_network(self):\n",
    "        \"\"\"Update the target Q network\"\"\"\n",
    "        self.target_dqn.set_weights(self.DQN.get_weights())\n",
    "\n",
    "    def add_experience(self, action, frame, reward, terminal, clip_reward=True):\n",
    "        \"\"\"Wrapper function for adding an experience to the Agent's replay buffer\"\"\"\n",
    "        self.replay_buffer.add_experience(action, frame, reward, terminal, clip_reward)\n",
    "\n",
    "    def learn(self, batch_size, gamma, frame_number, priority_scale=1.0):\n",
    "        \"\"\"Sample a batch and use it to improve the DQN\n",
    "        Arguments:\n",
    "            batch_size: How many samples to draw for an update\n",
    "            gamma: Reward discount\n",
    "            frame_number: Global frame number (used for calculating importances)\n",
    "            priority_scale: How much to weight priorities when sampling the replay buffer. 0 = completely random, 1 = completely based on priority\n",
    "        Returns:\n",
    "            The loss between the predicted and target Q as a float\n",
    "        \"\"\"\n",
    "\n",
    "        if self.use_per:\n",
    "            (states, actions, rewards, new_states, terminal_flags), importance, indices = self.replay_buffer.get_minibatch(batch_size=self.batch_size, priority_scale=priority_scale)\n",
    "            importance = importance ** (1-self.calc_epsilon(frame_number))\n",
    "        else:\n",
    "            states, actions, rewards, new_states, terminal_flags = self.replay_buffer.get_minibatch(batch_size=self.batch_size, priority_scale=priority_scale)\n",
    "\n",
    "        # Main DQN estimates best action in new states\n",
    "        arg_q_max = self.DQN.predict(new_states).argmax(axis=1)\n",
    "\n",
    "        # Target DQN estimates q-vals for new states\n",
    "        future_q_vals = self.target_dqn.predict(new_states)\n",
    "        double_q = future_q_vals[range(batch_size), arg_q_max]\n",
    "\n",
    "        # Calculate targets (bellman equation)\n",
    "        target_q = rewards + (gamma*double_q * (1-terminal_flags))\n",
    "\n",
    "        # Use targets to calculate loss (and use loss to calculate gradients)\n",
    "        with tf.GradientTape() as tape:\n",
    "            q_values = self.DQN(states)\n",
    "\n",
    "            one_hot_actions = tf.keras.utils.to_categorical(actions, self.n_actions, dtype=np.float32)  # using tf.one_hot causes strange errors\n",
    "            Q = tf.reduce_sum(tf.multiply(q_values, one_hot_actions), axis=1)\n",
    "\n",
    "            error = Q - target_q\n",
    "            loss = tf.keras.losses.Huber()(target_q, Q)\n",
    "\n",
    "            if self.use_per:\n",
    "                # Multiply the loss by importance, so that the gradient is also scaled.\n",
    "                # The importance scale reduces bias against situataions that are sampled\n",
    "                # more frequently.\n",
    "                loss = tf.reduce_mean(loss * importance)\n",
    "\n",
    "        model_gradients = tape.gradient(loss, self.DQN.trainable_variables)\n",
    "        self.DQN.optimizer.apply_gradients(zip(model_gradients, self.DQN.trainable_variables))\n",
    "\n",
    "        if self.use_per:\n",
    "            self.replay_buffer.set_priorities(indices, error)\n",
    "\n",
    "        return float(loss.numpy()), error\n",
    "\n",
    "    def save(self, folder_name, **kwargs):\n",
    "        \"\"\"Saves the Agent and all corresponding properties into a folder\n",
    "        Arguments:\n",
    "            folder_name: Folder in which to save the Agent\n",
    "            **kwargs: Agent.save will also save any keyword arguments passed.  This is used for saving the frame_number\n",
    "        \"\"\"\n",
    "\n",
    "        # Create the folder for saving the agent\n",
    "        if not os.path.isdir(folder_name):\n",
    "            os.makedirs(folder_name)\n",
    "\n",
    "        # Save DQN and target DQN\n",
    "        self.DQN.save(folder_name + '/dqn.h5')\n",
    "        self.target_dqn.save(folder_name + '/target_dqn.h5')\n",
    "\n",
    "        # Save replay buffer\n",
    "        self.replay_buffer.save(folder_name + '/replay-buffer')\n",
    "\n",
    "        # Save meta\n",
    "        with open(folder_name + '/meta.json', 'w+') as f:\n",
    "            f.write(json.dumps({**{'buff_count': self.replay_buffer.count, 'buff_curr': self.replay_buffer.current}, **kwargs}))  # save replay_buffer information and any other information\n",
    "\n",
    "    def load(self, folder_name, load_replay_buffer=True):\n",
    "        \"\"\"Load a previously saved Agent from a folder\n",
    "        Arguments:\n",
    "            folder_name: Folder from which to load the Agent\n",
    "        Returns:\n",
    "            All other saved attributes, e.g., frame number\n",
    "        \"\"\"\n",
    "\n",
    "        if not os.path.isdir(folder_name):\n",
    "            raise ValueError(f'{folder_name} is not a valid directory')\n",
    "\n",
    "        # Load DQNs\n",
    "        self.DQN = tf.keras.models.load_model(folder_name + '/dqn.h5')\n",
    "        self.target_dqn = tf.keras.models.load_model(folder_name + '/target_dqn.h5')\n",
    "        self.optimizer = self.DQN.optimizer\n",
    "\n",
    "        # Load replay buffer\n",
    "        if load_replay_buffer:\n",
    "            self.replay_buffer.load(folder_name + '/replay-buffer')\n",
    "\n",
    "        # Load meta\n",
    "        with open(folder_name + '/meta.json', 'r') as f:\n",
    "            meta = json.load(f)\n",
    "\n",
    "        if load_replay_buffer:\n",
    "            self.replay_buffer.count = meta['buff_count']\n",
    "            self.replay_buffer.current = meta['buff_curr']\n",
    "\n",
    "        del meta['buff_count'], meta['buff_curr']  # we don't want to return this information\n",
    "        return meta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create environment\n",
    "game_wrapper = GameWrapper(ENV_NAME, MAX_NOOP_STEPS)\n",
    "print(\"The environment has the following {} actions: {}\".format(game_wrapper.env.action_space.n, game_wrapper.env.unwrapped.get_action_meanings()))\n",
    "\n",
    "# TensorBoard writer\n",
    "writer = tf.summary.create_file_writer(TENSORBOARD_DIR)\n",
    "\n",
    "# Build main and target networks\n",
    "MAIN_DQN = build_q_network(game_wrapper.env.action_space.n, LEARNING_RATE, input_shape=INPUT_SHAPE)\n",
    "TARGET_DQN = build_q_network(game_wrapper.env.action_space.n, input_shape=INPUT_SHAPE)\n",
    "\n",
    "replay_buffer = ReplayBuffer(size=MEM_SIZE, input_shape=INPUT_SHAPE, use_per=USE_PER)\n",
    "agent = Agent(MAIN_DQN, TARGET_DQN, replay_buffer, game_wrapper.env.action_space.n, input_shape=INPUT_SHAPE, batch_size=BATCH_SIZE, use_per=USE_PER)\n",
    "\n",
    "# Training and evaluation\n",
    "if LOAD_FROM is None:\n",
    "    frame_number = 0\n",
    "    rewards = []\n",
    "    loss_list = []\n",
    "else:\n",
    "    print('Loading from', LOAD_FROM)\n",
    "    meta = agent.load(LOAD_FROM, LOAD_REPLAY_BUFFER)\n",
    "\n",
    "    # Apply information loaded from meta\n",
    "    frame_number = meta['frame_number']\n",
    "    rewards = meta['rewards']\n",
    "    loss_list = meta['loss_list']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Main loop\n",
    "try:\n",
    "    with writer.as_default():\n",
    "        while frame_number < TOTAL_FRAMES:\n",
    "            # Training\n",
    "\n",
    "            epoch_frame = 0\n",
    "            while epoch_frame < FRAMES_BETWEEN_EVAL:\n",
    "                start_time = time.time()\n",
    "                game_wrapper.reset()\n",
    "                life_lost = True\n",
    "                episode_reward_sum = 0\n",
    "                for _ in range(MAX_EPISODE_LENGTH):\n",
    "                    # Get action\n",
    "                    action = agent.get_action(frame_number, game_wrapper.state)\n",
    "\n",
    "                    # Take step\n",
    "                    processed_frame, reward, terminal, life_lost = game_wrapper.step(action)\n",
    "                    frame_number += 1\n",
    "                    epoch_frame += 1\n",
    "                    episode_reward_sum += reward\n",
    "\n",
    "                    # Add experience to replay memory\n",
    "                    agent.add_experience(action=action,\n",
    "                                        frame=processed_frame[:, :, 0],\n",
    "                                        reward=reward, clip_reward=CLIP_REWARD,\n",
    "                                        terminal=life_lost)\n",
    "\n",
    "                    # Update agent\n",
    "                    if frame_number % UPDATE_FREQ == 0 and agent.replay_buffer.count > MIN_REPLAY_BUFFER_SIZE:\n",
    "                        loss, _ = agent.learn(BATCH_SIZE, gamma=DISCOUNT_FACTOR, frame_number=frame_number, priority_scale=PRIORITY_SCALE)\n",
    "                        loss_list.append(loss)\n",
    "\n",
    "                    # Update target network\n",
    "                    if frame_number % UPDATE_FREQ == 0 and frame_number > MIN_REPLAY_BUFFER_SIZE:\n",
    "                        agent.update_target_network()\n",
    "\n",
    "                    # Break the loop when the game is over\n",
    "                    if terminal:\n",
    "                        terminal = False\n",
    "                        break\n",
    "\n",
    "                rewards.append(episode_reward_sum)\n",
    "\n",
    "                # Output the progress every 10 games\n",
    "                if len(rewards) % 10 == 0:\n",
    "                    # Write to TensorBoard\n",
    "                    if WRITE_TENSORBOARD:\n",
    "                        tf.summary.scalar('Reward', np.mean(rewards[-10:]), frame_number)\n",
    "                        tf.summary.scalar('Loss', np.mean(loss_list[-100:]), frame_number)\n",
    "                        writer.flush()\n",
    "\n",
    "                    print(f'Game number: {str(len(rewards)).zfill(6)}  Frame number: {str(frame_number).zfill(8)}  Average reward: {np.mean(rewards[-10:]):0.1f}  Time taken: {(time.time() - start_time):.1f}s')\n",
    "\n",
    "            # Evaluation every `FRAMES_BETWEEN_EVAL` frames\n",
    "            terminal = True\n",
    "            eval_rewards = []\n",
    "            evaluate_frame_number = 0\n",
    "\n",
    "            for _ in range(EVAL_LENGTH):\n",
    "                if terminal:\n",
    "                    game_wrapper.reset(evaluation=True)\n",
    "                    life_lost = True\n",
    "                    episode_reward_sum = 0\n",
    "                    terminal = False\n",
    "\n",
    "                # Breakout requires a \"fire\" action (action #1) to start the\n",
    "                # game each time a life is lost.\n",
    "                # Otherwise, the agent would sit around doing nothing.\n",
    "                action = 1 if life_lost else agent.get_action(frame_number, game_wrapper.state, evaluation=True)\n",
    "\n",
    "                # Step action\n",
    "                _, reward, terminal, life_lost = game_wrapper.step(action)\n",
    "                evaluate_frame_number += 1\n",
    "                episode_reward_sum += reward\n",
    "\n",
    "                # On game-over\n",
    "                if terminal:\n",
    "                    eval_rewards.append(episode_reward_sum)\n",
    "\n",
    "            if len(eval_rewards) > 0:\n",
    "                final_score = np.mean(eval_rewards)\n",
    "            else:\n",
    "                # In case the game is longer than the number of frames allowed\n",
    "                final_score = episode_reward_sum\n",
    "            # Print score and write to tensorboard\n",
    "            print('Evaluation score:', final_score)\n",
    "            if WRITE_TENSORBOARD:\n",
    "                tf.summary.scalar('Evaluation score', final_score, frame_number)\n",
    "                writer.flush()\n",
    "\n",
    "            # Save model\n",
    "            if len(rewards) > 300 and SAVE_PATH is not None:\n",
    "                agent.save(f'{SAVE_PATH}/save-{str(frame_number).zfill(8)}', frame_number=frame_number, rewards=rewards, loss_list=loss_list)\n",
    "except KeyboardInterrupt:\n",
    "    print('\\nTraining exited early.')\n",
    "    writer.close()\n",
    "\n",
    "    if SAVE_PATH is None:\n",
    "        try:\n",
    "            SAVE_PATH = input('Would you like to save the trained model? If so, type in a save path, otherwise, interrupt with ctrl+c. ')\n",
    "        except KeyboardInterrupt:\n",
    "            print('\\nExiting...')\n",
    "\n",
    "    if SAVE_PATH is not None:\n",
    "        print('Saving...')\n",
    "        agent.save(f'{SAVE_PATH}/save-{str(frame_number).zfill(8)}', frame_number=frame_number, rewards=rewards, loss_list=loss_list)\n",
    "        print('Saved.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# To watch the agent learn to play, head on over to evaluation.py and\n",
    "# change `RESTORE_PATH` to the location your model is saved."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
